import pysam
import argparse
import pod5 as p5
from tqdm import tqdm
from pathlib import Path

import pandas as pd
from collections import defaultdict
import openpyxl

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--tsv', type=str, default="/homeb/hpc/data/hg002.r10.4khz.rep1.chr2.4.tsv", required=False)
    parser.add_argument('--input', type=str, default="/home/xiaoyf/data/analyze/kmer_search.xlsx", required=False)
    parser.add_argument('--output', type=str, default="/home/xiaoyf/data/analyze/kmer_output.xlsx", required=False)
    return parser.parse_args()

def deal_bam(tsv_file, nested_dict, output_file):
    read_num = 0
    query_length = []
    new_nested_dict = defaultdict(lambda: defaultdict(list))
    
    # 读取 TSV 文件
    with open(tsv_file, 'r') as f:
        for line in f:
            words = line.strip().split('\t')
            read_name = words[0]
            reference_name = words[1]
            seq = words[3]
            bisulfite = [float(x) for x in words[15].split(",")]
            
            # 跳过不是 chr2 和 chr4 的行
            if reference_name != 'chr2' and reference_name != 'chr4':
                continue
            
            # 处理预测和参考位点的列表
            pred_pos = [int(x) for x in words[5].split(",")]
            ref_pos = [int(x) for x in words[-1].split(",")]
            
            # 遍历每个预测位点
            for i in range(len(pred_pos)):
                if ref_pos[i] == -1:
                    continue
                
                key = (reference_name, ref_pos[i])
                
                # 查找该 key 是否在嵌套字典中
                for sheet_name in nested_dict.keys():
                    if key in nested_dict[sheet_name].keys():
                        # 获取 rseq 序列
                        rseq = seq[pred_pos[i] - 20:pred_pos[i] + 21]
                        
                        # 将结果添加到 new_nested_dict 中
                        new_nested_dict[sheet_name][key].append([
                            nested_dict[sheet_name][key]['strand'],
                            nested_dict[sheet_name][key]['coverage'],
                            nested_dict[sheet_name][key]['rmet'],
                            nested_dict[sheet_name][key]['beizhu'],
                            rseq,
                            bisulfite[i]
                        ])
    
    # 写回 Excel 文件
    with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
        for sheet_name, data in new_nested_dict.items():
            records = []
            for key, values_list in data.items():
                chrom, pos = key
                for values in values_list:
                    records.append({
                        'chrom': chrom,
                        'pos': pos,
                        'strand': values[0],
                        'coverage': values[1],
                        'rmet': values[2],
                        'beizhu': values[3],
                        'rseq': values[4],
                        'bisulfite': values[5]
                    })
            
            # 将记录写入 DataFrame
            df_output = pd.DataFrame(records)
            # 写入当前表
            df_output.to_excel(writer, sheet_name=sheet_name, index=False)

    print(f"数据已处理并写入到 {output_file}")

if __name__ == '__main__':
    args = parse_args()
    
    # 读取输入 Excel 文件，构建 nested_dict
    sheets = pd.read_excel(args.input, sheet_name=None)
    nested_dict = {}

    for sheet_name, df in sheets.items():
        sheet_dict = {}
        for _, row in df.iterrows():
            key = (row['chrom'], int(row['pos']))
            sheet_dict[key] = {
                'strand': row['strand'],
                'coverage': row['coverage'],
                'rmet': row['rmet'],
                'beizhu': row['beizhu']
            }
        nested_dict[sheet_name] = sheet_dict
    
    print('开始处理 BAM 文件')
    deal_bam(args.tsv, nested_dict, args.output)
